import paramiko
import threading
import json
import time
import select
import re

class SSHClientWrapper:
    def __init__(self, hostname, port, username, password,window):
        self.ssh_client = paramiko.SSHClient()
        self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.hostname = hostname
        self.port = port
        self.username = username
        self.password = password
        self.channel = None
        self.last_used_time = time.time()
        self.window=window

    def connect(self):
        self.ssh_client.connect(hostname=self.hostname, port=self.port, username=self.username, password=self.password)
        self.channel = self.ssh_client.invoke_shell()
        self.read_thread = threading.Thread(target=self.read_output, args=(self.channel,))
        self.read_thread.daemon = True
        self.read_thread.start()
        self.last_used_time = time.time()

    def read_output(self, channel):
        ansi_escape = re.compile(r'\x1B[@-_][0-?]*[ -/]*[@-~]')
        while True:
            try:
                # 使用 select 模块非阻塞地检查通道是否准备好读取数据
                if select.select([channel], [], [], 0.1)[0]:
                    output = channel.recv(65535).decode('utf-8')
                    if not output:
                        break
                    clean_output = ansi_escape.sub('', output)
                    escaped_output = json.dumps(clean_output)[1:-1]  # 使用 json.dumps 转义特殊字符
                    # print("escaped_output:"+escaped_output)
                    self.window.evaluate_js(f"termSend(\"{escaped_output}\")")
            except paramiko.SSHException as e:
                print(f"SSHException: {e}")
                break
            except Exception as e:
                print(f"Exception: {e}")
                break
    def send_command(self, command):
        if self.channel:
            self.channel.send(command + "\n")
            self.last_used_time = time.time()

    def close(self):
        if self.channel:
            self.channel.close()
        self.ssh_client.close()

class SSHConnectionPool:
    def __init__(self, pool_size=5, idle_timeout=30, window=None):
        self.pool = {}
        self.pool_size = pool_size
        self.idle_timeout = idle_timeout
        self.lock = threading.Lock()
        self.window = window;

    def send_command(self,hostname,command):
        if "host_" + hostname not in self.pool:
            return ;
        ssh_wrapper = self.pool["host_" + hostname]
        if ssh_wrapper.channel:
            ssh_wrapper.channel.send(command + "\n")
            ssh_wrapper.last_used_time = time.time()
    def get_connection(self, hostname, port, username, password):
        if "host_"+hostname not in self.pool:
            self._initialize_pool(hostname, port, username, password)
        ssh_wrapper = self.pool["host_"+hostname]
        ssh_wrapper.last_used_time = time.time()
        return ssh_wrapper

    # def release_connection(self, ssh_wrapper):
    #     with self.lock:
    #         hostname = ssh_wrapper.hostname
    #         if hostname in self.pool:
    #             self.pool[hostname].put(ssh_wrapper)
    #             self._check_idle_connections(hostname)

    def _initialize_pool(self, hostname, port, username, password):
        print("_initialize_pool")
        # for _ in range(self.pool_size):
        ssh_wrapper = SSHClientWrapper(hostname=hostname, port=port, username=username, password=password,window=self.window)
        ssh_wrapper.connect()
        self.pool["host_"+hostname]=ssh_wrapper;

    def _check_idle_connections(self, hostname):
        with self.lock:
            if hostname in self.pool:
                current_time = time.time()
                while not self.pool[hostname].empty():
                    ssh_wrapper = self.pool[hostname].get()
                    if current_time - ssh_wrapper.last_used_time > self.idle_timeout:
                        ssh_wrapper.close()
                    else:
                        self.pool[hostname].put(ssh_wrapper)
                        break

    def close_all_connections(self):
        print("close_all_connections")
        # with self.lock:
        #     for hostname in list(self.pool.keys()):
        #         while not self.pool[hostname].empty():
        #             ssh_wrapper = self.pool[hostname].get()
        #             ssh_wrapper.close()
        #         del self.pool[hostname]

# 使用示例
if __name__ == "__main__":
    pool = SSHConnectionPool(pool_size=5, idle_timeout=30)

    # try:
    while True:
        command = input("$")  # 获取用户输入\
        if command.lower() == "exit":  # 输入 exit 退出
            break
        ssh_wrapper = pool.get_connection(hostname='1.94.47.246', port=22, username='root', password='dk2024!@%-_=+')
        ssh_wrapper.send_command(command)
        # pool.release_connection(ssh_wrapper)
    # finally:
    #     pool.close_all_connections()